// 多级记忆系统

use std::collections::HashMap;
use crate::symbol_table::SymbolTable;
use crate::utils::{is_common_word, calculate_semantic_similarity};

// 记忆类型枚举
#[derive(Debug, Clone, PartialEq)]
pub enum MemoryType {
    ShortTerm,   // 短期记忆
    MediumTerm,  // 中期记忆
    LongTerm,    // 长期记忆
}

// 记忆片段
#[derive(Debug, Clone)]
pub struct MemoryFragment {
    pub id: usize,
    pub content: Vec<usize>,  // 符号ID序列
    pub memory_type: MemoryType,
    pub importance: f32,
    pub created_at: u64,
    pub last_accessed: u64,
    pub access_count: u64,
    pub tags: Vec<String>,    // 标签
    pub related_fragments: Vec<usize>,  // 相关记忆片段ID
}

// 记忆系统
#[derive(Debug)]
pub struct MemorySystem {
    // 短期记忆（容量有限）
    pub short_term: Vec<MemoryFragment>,
    // 中期记忆（容量适中）
    pub medium_term: HashMap<usize, MemoryFragment>,
    // 长期记忆（容量大）
    pub long_term: HashMap<usize, MemoryFragment>,
    // 当前时间
    pub current_time: u64,
    // 下一个记忆片段ID
    pub next_fragment_id: usize,
    // 短期记忆容量
    pub short_term_capacity: usize,
    // 重要性阈值
    pub importance_threshold: f32,
}

impl MemorySystem {
    pub fn new() -> Self {
        Self {
            short_term: Vec::new(),
            medium_term: HashMap::new(),
            long_term: HashMap::new(),
            current_time: 0,
            next_fragment_id: 0,
            short_term_capacity: 7,  // 短期记忆容量为7个项目（基于米勒定律）
            importance_threshold: 0.5,  // 重要性阈值
        }
    }
    
    // 更新时间
    pub fn update_time(&mut self) {
        self.current_time += 1;
    }
    
    // 添加到短期记忆
    pub fn add_short_term(&mut self, content: Vec<usize>, tags: Vec<String>) -> usize {
        let fragment = MemoryFragment {
            id: self.next_fragment_id,
            content,
            memory_type: MemoryType::ShortTerm,
            importance: 0.0,  // 初始重要性为0
            created_at: self.current_time,
            last_accessed: self.current_time,
            access_count: 1,
            tags,
            related_fragments: Vec::new(),
        };
        
        let id = fragment.id;
        self.next_fragment_id += 1;
        
        // 添加到短期记忆
        self.short_term.push(fragment);
        
        // 如果超过容量，将最旧的转移到中期记忆
        if self.short_term.len() > self.short_term_capacity {
            if !self.short_term.is_empty() {
                let oldest = self.short_term.remove(0);
                self.extract_to_medium_term(oldest);
            }
        }
        
        id
    }
    
    // 提取到中期记忆
    pub fn extract_to_medium_term(&mut self, mut fragment: MemoryFragment) {
        // 计算重要性
        fragment.importance = self.calculate_importance(&fragment);
        
        // 更新记忆类型
        fragment.memory_type = MemoryType::MediumTerm;
        
        // 如果重要性足够高，直接转移到长期记忆
        if fragment.importance >= self.importance_threshold {
            self.consolidate_to_long_term(fragment);
        } else {
            // 添加到中期记忆
            self.medium_term.insert(fragment.id, fragment);
        }
    }
    
    // 计算记忆片段的重要性
    pub fn calculate_importance(&self, fragment: &MemoryFragment) -> f32 {
        let mut importance = 0.0;
        
        // 基于访问频率
        importance += fragment.access_count as f32 * 0.2;
        
        // 基于时间衰减
        let time_factor = 1.0 - (self.current_time - fragment.created_at) as f32 / 1000.0;
        importance += time_factor.max(0.0) * 0.3;
        
        // 基于内容长度
        importance += fragment.content.len() as f32 * 0.1;
        
        // 基于标签数量
        importance += fragment.tags.len() as f32 * 0.1;
        
        // 基于相关片段数量
        importance += fragment.related_fragments.len() as f32 * 0.1;
        
        // 基于内容中的关键词
        // 这里可以添加更复杂的关键词检测逻辑
        
        importance
    }
    
    // 提取主题
    pub fn extract_topics(&self, symbol_table: &SymbolTable, top_n: usize) -> Vec<usize> {
        let mut symbol_frequency = HashMap::new();
        
        // 统计所有记忆片段中的符号频率
        for fragment in self.short_term.iter().chain(self.medium_term.values()).chain(self.long_term.values()) {
            for &symbol_id in &fragment.content {
                *symbol_frequency.entry(symbol_id).or_insert(0) += 1;
            }
        }
        
        // 按频率排序
        let mut symbols: Vec<_> = symbol_frequency.into_iter().collect();
        symbols.sort_by(|a, b| b.1.cmp(&a.1));
        
        // 过滤常见词并返回前top_n个
        symbols.into_iter()
            .filter_map(|(id, _)| {
                if let Some(text) = symbol_table.get_text(id) {
                    if !is_common_word(text) {
                        Some(id)
                    } else {
                        None
                    }
                } else {
                    None
                }
            })
            .take(top_n)
            .collect()
    }
    
    // 搜索相关记忆
    pub fn search_relevant_memories(&self, query: &[usize], symbol_table: &SymbolTable, max_results: usize) -> Vec<&MemoryFragment> {
        let mut results = Vec::new();
        
        // 搜索短期记忆
        for fragment in &self.short_term {
            let relevance = self.calculate_relevance(query, &fragment.content, symbol_table);
            if relevance > 0.1 {  // 设置最小相关性阈值
                results.push((fragment, relevance));
            }
        }
        
        // 搜索中期记忆
        for fragment in self.medium_term.values() {
            let relevance = self.calculate_relevance(query, &fragment.content, symbol_table);
            if relevance > 0.1 {
                results.push((fragment, relevance));
            }
        }
        
        // 搜索长期记忆
        for fragment in self.long_term.values() {
            let relevance = self.calculate_relevance(query, &fragment.content, symbol_table);
            if relevance > 0.1 {
                results.push((fragment, relevance));
            }
        }
        
        // 按相关性排序
        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        
        // 返回前max_results个结果
        results.into_iter().take(max_results).map(|(fragment, _)| fragment).collect()
    }
    
    // 计算查询与记忆片段的相关性
    fn calculate_relevance(&self, query: &[usize], content: &[usize], symbol_table: &SymbolTable) -> f32 {
        if query.is_empty() || content.is_empty() {
            return 0.0;
        }
        
        // 简化相关性计算，避免创建大型HashMap
        let mut common_symbols = 0;
        for &sym1 in query {
            for &sym2 in content {
                if sym1 == sym2 {
                    common_symbols += 1;
                    break;
                }
            }
        }
        
        // 计算直接匹配的相关性
        let direct_match = common_symbols as f32 / query.len().max(content.len()) as f32;
        
        // 如果有直接匹配，返回较高的相关性
        if direct_match > 0.0 {
            return direct_match;
        }
        
        // 否则返回低相关性
        0.1
    }
    
    // 巩固到长期记忆
    pub fn consolidate_to_long_term(&mut self, mut fragment: MemoryFragment) {
        // 更新记忆类型
        fragment.memory_type = MemoryType::LongTerm;
        
        // 添加到长期记忆
        self.long_term.insert(fragment.id, fragment);
    }
    
    // 访问记忆片段
    pub fn access_memory(&mut self, fragment_id: usize) -> Option<&MemoryFragment> {
        // 更新时间
        self.update_time();
        
        // 首先检查短期记忆
        for i in 0..self.short_term.len() {
            if self.short_term[i].id == fragment_id {
                self.short_term[i].last_accessed = self.current_time;
                self.short_term[i].access_count += 1;
                return Some(&self.short_term[i]);
            }
        }
        
        // 检查中期记忆
        if let Some(fragment) = self.medium_term.get(&fragment_id) {
            let fragment_id = fragment.id;
            let last_accessed = self.current_time;
            let access_count = fragment.access_count + 1;
            let importance = fragment.importance;
            
            // 如果访问频率足够高，考虑提升到长期记忆
            if importance >= self.importance_threshold {
                let fragment = fragment.clone();
                self.medium_term.remove(&fragment_id);
                self.consolidate_to_long_term(fragment);
                return self.long_term.get(&fragment_id);
            }
            
            // 更新访问信息
            if let Some(fragment) = self.medium_term.get_mut(&fragment_id) {
                fragment.last_accessed = last_accessed;
                fragment.access_count = access_count;
            }
            
            return self.medium_term.get(&fragment_id);
        }
        
        // 检查长期记忆
        if let Some(fragment) = self.long_term.get_mut(&fragment_id) {
            fragment.last_accessed = self.current_time;
            fragment.access_count += 1;
            return Some(fragment);
        }
        
        None
    }
    
    // 应用时间衰减
    pub fn apply_time_decay(&mut self, decay_factor: f32) {
        // 对中期记忆应用时间衰减
        for fragment in self.medium_term.values_mut() {
            fragment.importance *= decay_factor;
        }
        
        // 对长期记忆应用时间衰减
        for fragment in self.long_term.values_mut() {
            fragment.importance *= decay_factor;
        }
    }
    
    // 定期整理中期记忆
    pub fn organize_medium_term(&mut self) {
        let mut to_promote = Vec::new();
        let mut to_remove = Vec::new();
        
        for (&id, fragment) in &self.medium_term {
            if fragment.importance >= self.importance_threshold {
                to_promote.push(id);
            } else if fragment.importance < 0.1 {
                to_remove.push(id);
            }
        }
        
        // 提升重要记忆到长期记忆
        for id in to_promote {
            if let Some(fragment) = self.medium_term.remove(&id) {
                self.consolidate_to_long_term(fragment);
            }
        }
        
        // 移除不重要的记忆
        for id in to_remove {
            self.medium_term.remove(&id);
        }
    }
    
    // 获取记忆统计信息
    pub fn get_stats(&self) -> (usize, usize, usize) {
        (
            self.short_term.len(),
            self.medium_term.len(),
            self.long_term.len(),
        )
    }
}